#!/usr/bin/env python3

import os
import mmap
import struct
import time
import numpy as np
import cv2
import rospy
import threading
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

class SharedMemoryManager:
    """修复后的共享内存管理器"""
    
    def __init__(self, shm_name="/image_shm", max_size=50 * 1024 * 1024):  # 50MB
        self.shm_name = shm_name
        self.max_size = max_size
        self.lock = threading.Lock()
        self._create_shared_memory()
        self.bridge = CvBridge()
    
    def _create_shared_memory(self):
        """创建共享内存区域"""
        try:
            os.unlink(self.shm_name)
        except FileNotFoundError:
            pass
        
        # 创建共享内存文件
        self.fd = os.open(self.shm_name, os.O_CREAT | os.O_RDWR, 0o666)
        os.ftruncate(self.fd, self.max_size)
        
        # 映射共享内存
        self.shm = mmap.mmap(self.fd, self.max_size, mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)
    
    def write_images(self, color_msg, depth_msg):
        """修复形状参数不匹配问题的写入方法，并记录每部分耗时"""
        with self.lock:
            try:
                # 记录总开始时间
                start_total = time.perf_counter()

                # 1. 转换彩色图像
                start_color = time.perf_counter()
                color_img = self.bridge.imgmsg_to_cv2(color_msg, "bgr8")
                color_time = time.perf_counter() - start_color
                rospy.loginfo(f"彩色图像转换耗时: {color_time:.4f} 秒")

                # 2. 转换和处理深度图像
                start_depth = time.perf_counter()
                depth_img = self.bridge.imgmsg_to_cv2(depth_msg, "passthrough")
                if depth_img.ndim == 2:
                    depth_img = depth_img[:, :, np.newaxis]  # 添加通道维度
                depth_time = time.perf_counter() - start_depth
                rospy.loginfo(f"深度图像转换和处理耗时: {depth_time:.4f} 秒")

                # 统一形状为 (H, W, C)
                color_shape = color_img.shape
                depth_shape = depth_img.shape

                # 3. 写入共享内存
                start_write = time.perf_counter()
                self.shm.seek(0)
                self.shm.write(b'VALID')
                self.shm.write(struct.pack('Q', int(time.time() * 1000)))  # 时间戳
                
                # 写入数据大小
                color_data = color_img.tobytes()
                depth_data = depth_img.tobytes()
                self.shm.write(struct.pack('II', len(color_data), len(depth_data)))
                
                # 写入形状 (H, W, C)
                self.shm.write(struct.pack('III', *color_shape))
                self.shm.write(struct.pack('III', *depth_shape))
                
                # 写入图像数据
                self.shm.write(color_data)
                self.shm.write(depth_data)
                write_time = time.perf_counter() - start_write
                rospy.loginfo(f"共享内存写入耗时: {write_time:.4f} 秒")

                # 总耗时
                total_time = time.perf_counter() - start_total
                rospy.loginfo(f"总写入耗时: {total_time:.4f} 秒")

                return True
            except Exception as e:
                rospy.logerr(f"写入共享内存错误: {e}")
                return False
    
    def close(self):
        """关闭共享内存"""
        try:
            self.shm.close()
            os.close(self.fd)
            os.unlink(self.shm_name)
        except:
            pass

class ImageAcquisition:
    def __init__(self):
        self.bridge = CvBridge()
        self.shm_manager = SharedMemoryManager()
        self.color_img = None
        self.depth_img = None
        self.lock = threading.Lock()
        
        # ROS订阅
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_callback)
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)
        rospy.loginfo("ImageAcquisition initialized, waiting for images...")
        
    def color_callback(self, msg):
        try:
            with self.lock:
                self.color_img = msg
                self.try_process_images()
        except Exception as e:
            rospy.logerr(f"Color error: {e}")

    def depth_callback(self, msg):
        try:
            with self.lock:
                self.depth_img = msg
                self.try_process_images()
        except Exception as e:
            rospy.logerr(f"Depth error: {e}")

    def try_process_images(self):
        if self.color_img is not None and self.depth_img is not None:
            success = self.shm_manager.write_images(self.color_img, self.depth_img)
            if success:
                rospy.loginfo_throttle(5, "Image pair written to shared memory")
            
            # 重置
            self.color_img = None
            self.depth_img = None

    def shutdown(self):
        """关闭时清理"""
        self.shm_manager.close()
        rospy.loginfo("ImageAcquisition shutdown")

if __name__ == "__main__":
    rospy.init_node("image_acquisition")
    acquisition = ImageAcquisition()
    rospy.on_shutdown(acquisition.shutdown)
    rospy.spin()